package com.zhoug.silerovad;

import android.content.Context;
import android.util.Log;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.nio.ByteBuffer;
import java.nio.ByteOrder;
import java.nio.FloatBuffer;
import java.nio.LongBuffer;
import java.nio.ShortBuffer;
import java.util.HashMap;
import java.util.Map;

import ai.onnxruntime.OnnxTensor;
import ai.onnxruntime.OrtEnvironment;
import ai.onnxruntime.OrtException;
import ai.onnxruntime.OrtSession;

/**
 * @description: 类的简要描述
 * @author: 35574
 * @date: 2025/10/27 9:17
 */
public class SileroVAD {
    public static final int STATE_UNINITIALIZED = 0;
    public static final int STATE_INITIALIZED = 1;
    private static final boolean DEBUG = false;
    private static final String TAG = ">>>SileroVAD";
    private int mState = STATE_UNINITIALIZED;
    private OrtEnvironment environment;
    private OrtSession session;
    // 模型期望的输入名称
    private static final String INPUT_NAME = "input";     // 音频数据
    private static final String STATE_NAME = "state";     // 模型状态
    private static final String SR_NAME = "sr";           // 采样率
    // 模型参数 支持 8000 Hz and 16000 Hz
    private int sampleRate = 16000;  // Silero VAD 要求的采样率
    private static final int EXPECTED_FRAME_SIZE = 512; // 模型期望的帧大小 即音频每次读取byte[]为1024
    // 模型状态
    private float[][][] vadState; // 形状为 [2, 1, 128]

    public void setSampleRate(int sampleRate) {
        this.sampleRate = sampleRate;
    }


    /**
     * 加载模型
     * @param context
     */
    public void loadModel(Context context) {
        try {
            environment = OrtEnvironment.getEnvironment();
            // 从assets读取模型文件
            InputStream modelStream = context.getAssets().open("silero_vad.onnx");
            byte[] modeBytes = readAllBytes(modelStream);
            OrtSession.SessionOptions sessionOptions = new OrtSession.SessionOptions();
            sessionOptions.setOptimizationLevel(OrtSession.SessionOptions.OptLevel.BASIC_OPT);
            session = environment.createSession(modeBytes, sessionOptions);
            // 打印模型输入信息，帮助调试
            // input, value: NodeInfo(name=input,info=TensorInfo(javaType=FLOAT,onnxType=ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT,shape=[-1, -1]))
            // state, value: NodeInfo(name=state,info=TensorInfo(javaType=FLOAT,onnxType=ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT,shape=[2, -1, 128]))
            // sr, value: NodeInfo(name=sr,info=TensorInfo(javaType=INT64,onnxType=ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64,shape=[]))
            //03.276 ActivityThread       com....audiorecorder  D  add activity client record, r= ActivityRecord{a0679b token=android.os.BinderProxy@c14ce84 {com.zhoug.demo.audiorecorder/com.zhoug.demo.audiorecorder
            /*Map<String, NodeInfo> inputInfo = session.getInputInfo();
            for (Map.Entry<String, NodeInfo> entry : inputInfo.entrySet()) {
                NodeInfo value = entry.getValue();
                logd("Input name: " + entry.getKey() + ", value: " + value.toString());
            }*/
            initializeState();
            mState = STATE_INITIALIZED;
        } catch (Exception e) {
            loge("loadModel: " + e.getMessage());
        }
    }

    /**
     * 初始化模型状态
     */
    private void initializeState() {
        // state 的形状是 [2, 1, 128]
        vadState = new float[2][1][128];
        // 初始状态设为0
        for (int i = 0; i < 2; i++) {
            for (int j = 0; j < 1; j++) {
                for (int k = 0; k < 128; k++) {
                    vadState[i][j][k] = 0.0f;
                }
            }
        }
    }

    /**
     * VAD 检测
     * @param audioData 采集的音频每帧pcm数据
     * @return
     */
    public float predict(byte[] audioData) {
        if (mState == STATE_UNINITIALIZED) {
            Log.e(TAG, "predict: 未初始化");
            return 0;
        }
        // 将16位PCM byte数组转换为short数组
        short[] shortData = byteToShortArray(audioData);
        // 将short数组转换为float数组，并归一化到[-1, 1]范围
        float[] floatData = shortToFloatArray(shortData);
        // 创建模型输入tensor
        // 创建音频数据tensor
        try {
            //打印日志发现 InputNames:input state sr
           /*  Set<String> inputNames = session.getInputNames();
           Iterator<String> iterator = inputNames.iterator();
            String inputName = null;
            while (iterator.hasNext()) {
                inputName = iterator.next();
               logd( "predict:inputName=" + inputName);
            }*/
            // 1. 准备输入 Tensor
            long[] inputShape = new long[]{1, (long) floatData.length};
            OnnxTensor inputTensor = OnnxTensor.createTensor(environment, FloatBuffer.wrap(floatData), inputShape);
            // 2. 准备状态输入数据
            long[] stateShape = {2, 1, 128};
            FloatBuffer stateBuffer = FloatBuffer.allocate(2 * 1 * 128);
            for (int i = 0; i < 2; i++) {
                for (int j = 0; j < 1; j++) {
                    for (int k = 0; k < 128; k++) {
                        stateBuffer.put(vadState[i][j][k]);
                    }
                }
            }
            stateBuffer.rewind();
            OnnxTensor stateTensor = OnnxTensor.createTensor(environment, stateBuffer, stateShape);

            // 3. 准备采样率输入数据
            long[] srShape = {1};
            LongBuffer srBuffer = LongBuffer.wrap(new long[]{sampleRate});
            OnnxTensor srTensor = OnnxTensor.createTensor(environment, srBuffer, srShape);
            // 4. 构建输入映射
            Map<String, OnnxTensor> inputs = new HashMap<>();
            inputs.put(INPUT_NAME, inputTensor);
            inputs.put(STATE_NAME, stateTensor);
            inputs.put(SR_NAME, srTensor);

            // 5. 运行模型推理
            OrtSession.Result results = session.run(inputs);

            // 6. 处理输出
            // 输出有两个：语音概率和新的状态
            float[][] output = (float[][]) results.get(0).getValue(); // 语音概率
            float[][][] newState = (float[][][]) results.get(1).getValue(); // 新的状态

            float speechProbability = output[0][0]; // 获取语音概率

            // 7. 更新状态
            this.vadState = newState;

            // 8. 清理资源
            inputTensor.close();
            stateTensor.close();
            srTensor.close();
            results.close();
            return speechProbability;

        } catch (Exception e) {
            loge("predict: " + e.getMessage());
            e.printStackTrace();
        }


        return 0f;
    }

    /**
     * 释放资源 ,需要重新{@link #loadModel(Context)}
     */
    public void release() {
        if (session != null) {
            try {
                session.close();
            } catch (OrtException e) {
                loge(e.getMessage());
            }
            session = null;
        }
        if (environment != null) {
            environment.close();
            environment = null;
        }
        mState = STATE_UNINITIALIZED;
    }

    /**
     * 获取状态
     *  {@link #STATE_UNINITIALIZED,#STATE_INITIALIZED}
     * @return
     */
    public int getState() {
        return mState;
    }


    /**
     * 将byte数组转换为short数组
     * 假设是16位小端序PCM数据
     */
    private static short[] byteToShortArray(byte[] byteData) {
        if (byteData == null || byteData.length % 2 != 0) {
            throw new IllegalArgumentException("Invalid byte array");
        }
        ByteBuffer buffer = ByteBuffer.wrap(byteData);
        buffer.order(ByteOrder.LITTLE_ENDIAN); // PCM数据通常是小端序
        ShortBuffer shortBuffer = buffer.asShortBuffer();
        short[] shortData = new short[byteData.length / 2];
        shortBuffer.get(shortData);
        return shortData;
    }

    /**
     * 将short数组转换为float数组
     * @param shortData
     * @return
     */
    private static float[] shortToFloatArray(short[] shortData) {
        // 将short数组转换为float数组，并归一化到[-1, 1]范围
        float[] floatData = new float[shortData.length];
        for (int i = 0; i < shortData.length; i++) {
            floatData[i] = shortData[i] / 32768.0f;
        }
        return floatData;
    }

    /**
     * 读取流中的数据到byte数组
     * @param ins
     * @return
     * @throws IOException
     */
    public static byte[] readAllBytes(InputStream ins) throws IOException {
        if (android.os.Build.VERSION.SDK_INT >= android.os.Build.VERSION_CODES.TIRAMISU) {
            return ins.readAllBytes();
        } else {
            ByteArrayOutputStream bos = new ByteArrayOutputStream();
            byte[] buffer = new byte[1024];
            int len;
            while ((len = ins.read(buffer)) != -1) {
                bos.write(buffer, 0, len);
            }
            ins.close();
            bos.close();
            return bos.toByteArray();
        }
    }


    private void logd(String msg) {
        if (DEBUG) {
            Log.d(TAG, msg);
        }
    }


    private void loge(String error) {
        Log.e(TAG, error);
    }

}
